sae explore

import%20marimo%0A%0A__generated_with%20%3D%20%220.18.4%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22%2C%20auto_download%3D%5B%22html%22%5D)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Imports%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20return%20mo%2C%20np%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20model.dinov3_sae_topk_model%20import%20Dinov3TopKSAEDPTSegmentation%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torchvision.transforms%20as%20T%0A%20%20%20%20from%20torchvision.transforms%20import%20v2%0A%20%20%20%20import%20rasterio%20as%20rio%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20from%20matplotlib.patches%20import%20Rectangle%0A%20%20%20%20import%20cv2%0A%20%20%20%20import%20plotly.graph_objects%20as%20go%0A%20%20%20%20from%20plotly.subplots%20import%20make_subplots%0A%20%20%20%20import%20torch.nn.functional%20as%20F%0A%20%20%20%20from%20plotly%20import%20express%20as%20px%0A%20%20%20%20import%20altair%20as%20alt%0A%20%20%20%20import%20pandas%20as%20pd%0A%20%20%20%20import%20base64%0A%20%20%20%20from%20io%20import%20BytesIO%0A%20%20%20%20from%20PIL%20import%20Image%0A%20%20%20%20import%20os%0A%20%20%20%20%23%20Set%20font%20to%20Times%20for%20ICML%202026%20compatibility%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20Dinov3TopKSAEDPTSegmentation%2C%0A%20%20%20%20%20%20%20%20F%2C%0A%20%20%20%20%20%20%20%20cv2%2C%0A%20%20%20%20%20%20%20%20go%2C%0A%20%20%20%20%20%20%20%20make_subplots%2C%0A%20%20%20%20%20%20%20%20os%2C%0A%20%20%20%20%20%20%20%20plt%2C%0A%20%20%20%20%20%20%20%20px%2C%0A%20%20%20%20%20%20%20%20rio%2C%0A%20%20%20%20%20%20%20%20torch%2C%0A%20%20%20%20%20%20%20%20v2%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(os)%3A%0A%20%20%20%20os.environ%5B'MARIMO_OUTPUT_MAX_BYTES'%5D%3D'20_000_000'%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(plt)%3A%0A%20%20%20%20plt.rcParams%5B'font.family'%5D%20%3D%20'serif'%0A%20%20%20%20plt.rcParams%5B'font.serif'%5D%20%3D%20%5B'Times%20New%20Roman'%2C%20'Times'%2C%20'DejaVu%20Serif'%5D%0A%20%20%20%20plt.rcParams%5B'mathtext.fontset'%5D%20%3D%20'cm'%20%20%23%20Computer%20Modern%20for%20math%0A%20%20%20%20plt.rcParams%5B'pdf.fonttype'%5D%20%3D%2042%20%20%23%20TrueType%20(Type-1)%20fonts%20in%20PDF%0A%20%20%20%20plt.rcParams%5B'ps.fonttype'%5D%20%3D%2042%20%20%20%23%20TrueType%20(Type-1)%20fonts%20in%20PS%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20chkpt_path%20%3D%20%22%2Fhome%2Fprod-gpu-3%2FDocuments%2Fth%2Ftraining_artifacts%2Fcheckpoints%2Fexperiment_20260119_121235%2Flast.ckpt%22%0A%20%20%20%20sae_hidden_dim%3A%20int%20%3D%2065536%0A%20%20%20%20num_classes%20%3D%209%0A%20%20%20%20device%20%3D%20%22cuda%3A0%22%0A%20%20%20%20img_size%20%3D%201024%0A%20%20%20%20patch_size%20%3D%2016%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20chkpt_path%2C%0A%20%20%20%20%20%20%20%20device%2C%0A%20%20%20%20%20%20%20%20img_size%2C%0A%20%20%20%20%20%20%20%20num_classes%2C%0A%20%20%20%20%20%20%20%20patch_size%2C%0A%20%20%20%20%20%20%20%20sae_hidden_dim%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Dinov3TopKSAEDPTSegmentation%2C%0A%20%20%20%20chkpt_path%2C%0A%20%20%20%20device%2C%0A%20%20%20%20img_size%2C%0A%20%20%20%20num_classes%2C%0A%20%20%20%20sae_hidden_dim%3A%20int%2C%0A%20%20%20%20torch%2C%0A%20%20%20%20v2%2C%0A)%3A%0A%20%20%20%20checkpoint%20%3D%20torch.load(chkpt_path%2C%20map_location%3Ddevice)%0A%20%20%20%20model%20%3D%20Dinov3TopKSAEDPTSegmentation(num_classes%20%3D%20num_classes%2C%20sae_hidden_dim%3Dsae_hidden_dim).to(device)%0A%20%20%20%20if%20'state_dict'%20in%20checkpoint%3A%0A%20%20%20%20%20%20%20%20state_dict%20%3D%20checkpoint%5B'state_dict'%5D%0A%20%20%20%20%20%20%20%20%23%20Remove%20'model.'%20prefix%20if%20present%0A%20%20%20%20%20%20%20%20state_dict%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20k.replace('model.'%2C%20'')%3A%20v%20%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20k%2C%20v%20in%20state_dict.items()%0A%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20model.load_state_dict(state_dict%2C%20strict%20%3D%20False)%0A%20%20%20%20model.eval()%0A%20%20%20%20model.inference_mode()%0A%20%20%20%20image_transform%20%3D%20v2.Compose(%5B%0A%20%20%20%20%20%20%20%20v2.Resize((img_size%2C%20img_size))%2C%0A%20%20%20%20%20%20%20%20v2.ToDtype(torch.float32%2C%20scale%20%3D%20True)%2C%0A%20%20%20%20%20%20%20%20v2.Normalize(mean%3D%5B0.485%2C%200.456%2C%200.406%5D%2C%20std%3D%5B0.229%2C%200.224%2C%200.225%5D)%0A%20%20%20%20%5D)%0A%20%20%20%20return%20image_transform%2C%20model%0A%0A%0A%40app.cell%0Adef%20_(plt%2C%20rio)%3A%0A%20%20%20%20img%20%3D%20None%0A%20%20%20%20with%20rio.open(%22data%2Flulc%2Fimages%2Faustin_45.tif%22)%20as%20src%3A%0A%20%20%20%20%20%20%20%20img%20%3D%20src.read()%0A%20%20%20%20plt.imshow(img.transpose(%5B1%2C2%2C0%5D))%0A%20%20%20%20return%20(img%2C)%0A%0A%0A%40app.cell%0Adef%20_(cv2%2C%20device%2C%20image_transform%2C%20img%2C%20img_size%2C%20torch)%3A%0A%20%20%20%20img_tensor%20%3D%20image_transform(torch.from_numpy(img)).unsqueeze(0).to(device)%0A%20%20%20%20img_rsz%20%3D%20cv2.resize(img.transpose(%5B1%2C2%2C0%5D)%2C%20(img_size%2C%20img_size)%2C%20interpolation%3Dcv2.INTER_LINEAR).transpose(%5B2%2C0%2C1%5D)%0A%20%20%20%20img_rsz.max()%0A%20%20%20%20return%20img_rsz%2C%20img_tensor%0A%0A%0A%40app.cell%0Adef%20_(model)%3A%0A%20%20%20%20backbone%20%3D%20model.backbone%0A%20%20%20%20sae%20%3D%20model.sae%0A%20%20%20%20head%20%3D%20model.segmentation_head%0A%20%20%20%20return%20backbone%2C%20head%2C%20sae%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%20Model%20Run%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20backbone%2C%20head%2C%20img_tensor%2C%20model%2C%20patch_size%2C%20sae%2C%20torch)%3A%0A%20%20%20%20inp%20%3D%20img_tensor%0A%20%20%20%20with%20torch.no_grad()%3A%0A%20%20%20%20%20%20%20%20B%2C%20C%2C%20H%2C%20W%20%3D%20inp.shape%0A%0A%20%20%20%20%20%20%20%20if%20H%20%25%20patch_size%20!%3D%200%20or%20W%20%25%20patch_size%20!%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20raise%20ValueError(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22Image%20size%20%7BH%7Dx%7BW%7D%20not%20divisible%20by%20patch%20size%20%7Bpatch_size%7D%22%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20patch_h%2C%20patch_w%20%3D%20H%20%2F%2F%20patch_size%2C%20W%20%2F%2F%20patch_size%0A%20%20%20%20%20%20%20%20feats%20%3D%20backbone(inp)%0A%20%20%20%20%20%20%20%20layer_features%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20cls_tokens%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20feat%2C%20cls_token%20in%20feats%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20layer_features.append(feat)%0A%20%20%20%20%20%20%20%20%20%20%20%20cls_tokens.append(cls_token)%0A%20%20%20%20%20%20%20%20sae_out%20%3D%20sae(layer_features)%0A%20%20%20%20%20%20%20%20layer_reconstructed%20%3D%20sae_out%5B'layer_reconstructed'%5D%0A%0A%20%20%20%20%20%20%20%20processed_features%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20i%2C%20recon%20in%20enumerate(layer_reconstructed)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20processed_features.append((recon%2C%20cls_tokens%5Bi%5D))%0A%0A%20%20%20%20%20%20%20%20logits%20%3D%20head(processed_features%2C%20patch_h%2C%20patch_w)%0A%0A%20%20%20%20class%20ModelSvs%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20model)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.backbone%20%3D%20model.backbone%0A%20%20%20%20%20%20%20%20%20%20%20%20self.sae%20%3D%20model.sae%0A%20%20%20%20%20%20%20%20%20%20%20%20self.W_dec%20%3D%20self.sae.W_dec%0A%20%20%20%20%20%20%20%20%20%20%20%20self.b_dec%20%3D%20self.sae.b_dec%0A%20%20%20%20%20%20%20%20%20%20%20%20self.head%20%3D%20model.segmentation_head%0A%0A%20%20%20%20%20%20%20%20def%20get_raw_feats(self%2C%20inp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20B%2C%20C%2C%20H%2C%20W%20%3D%20inp.shape%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20H%20%25%20patch_size%20!%3D%200%20or%20W%20%25%20patch_size%20!%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20raise%20ValueError(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22Image%20size%20%7BH%7Dx%7BW%7D%20not%20divisible%20by%20patch%20size%20%7Bpatch_size%7D%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20patch_h%2C%20patch_w%20%3D%20H%20%2F%2F%20patch_size%2C%20W%20%2F%2F%20patch_size%0A%20%20%20%20%20%20%20%20%20%20%20%20feats%20%3D%20self.backbone(inp)%0A%20%20%20%20%20%20%20%20%20%20%20%20layer_features%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20cls_tokens%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20feat%2C%20cls_token%20in%20feats%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20layer_features.append(feat)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cls_tokens.append(cls_token)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20layer_features%2C%20cls_tokens%0A%0A%20%20%20%20%20%20%20%20def%20predict_custom_features(self%2C%20inp%2C%20w_dec%2C%20b_dec)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20z%2C%20cls_tokens%20%3D%20self.get_raw_feats(inp)%0A%20%20%20%20%20%20%20%20%20%20%20%20h%2C%20z_concat%20%3D%20self.sae.encode(z%2C%20apply_topk%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20z_reconstructed%20%3D%20F.linear(h%2C%20w_dec)%20%2B%20b_dec%0A%20%20%20%20%20%20%20%20%20%20%20%20layer_reconstructed%20%3D%20self.sae.split_layers(z_reconstructed)%0A%20%20%20%20%20%20%20%20%20%20%20%20logits%20%3D%20self.get_logits(layer_reconstructed%2C%20cls_tokens)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20logits.argmax(dim%20%3D%201)%0A%0A%0A%20%20%20%20%20%20%20%20%40torch.no_grad%0A%20%20%20%20%20%20%20%20def%20get_sparse_coeffs(self%2C%20inp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20feats%2C%20cls_tokens%20%3D%20self.get_raw_feats(inp)%0A%20%20%20%20%20%20%20%20%20%20%20%20h%2C%20z_concat%20%3D%20self.sae.encode(feats%2C%20apply_topk%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20h%2C%20cls_tokens%0A%0A%20%20%20%20%20%20%20%20%40torch.no_grad%0A%20%20%20%20%20%20%20%20def%20get_output_from_sparse_ceoff(self%2C%20h%2C%20cls_tokens)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20z_reconstructed%20%3D%20self.sae.decode(h)%0A%20%20%20%20%20%20%20%20%20%20%20%20layer_reconstructed%20%3D%20self.sae.split_layers(z_reconstructed)%0A%20%20%20%20%20%20%20%20%20%20%20%20logits%20%3D%20self.get_logits(layer_reconstructed%2C%20cls_tokens)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20logits.argmax(dim%20%3D%201)%0A%0A%0A%20%20%20%20%20%20%20%20def%20get_sae(self%2C%20z)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20h%2C%20z_concat%20%3D%20self.sae.encode(z%2C%20apply_topk%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20h_flat%20%3D%20h%0A%20%20%20%20%20%20%20%20%20%20%20%20z_reconstructed%20%3D%20self.sae.decode(h)%0A%20%20%20%20%20%20%20%20%20%20%20%20layer_reconstructed%20%3D%20self.sae.split_layers(z_reconstructed)%0A%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20layer_reconstructed%0A%0A%20%20%20%20%20%20%20%20def%20get_logits(self%2C%20layer_reconstructed%2C%20cls_tokens)%3A%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20processed_features%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20i%2C%20recon%20in%20enumerate(layer_reconstructed)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20processed_features.append((recon%2C%20cls_tokens%5Bi%5D))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20logits%20%3D%20self.head(processed_features%2C%20patch_h%2C%20patch_w)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20logits%0A%0A%20%20%20%20%20%20%20%20%40torch.no_grad%0A%20%20%20%20%20%20%20%20def%20predict(self%2C%20inp)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20feats%2C%20cls_tokens%20%3D%20self.get_raw_feats(inp)%0A%20%20%20%20%20%20%20%20%20%20%20%20layer_recon%20%3D%20self.get_sae(feats)%0A%20%20%20%20%20%20%20%20%20%20%20%20logits%20%3D%20self.get_logits(layer_recon%2C%20cls_tokens)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20logits.argmax(dim%20%3D%201)%0A%0A%20%20%20%20model_service%20%3D%20ModelSvs(model)%0A%20%20%20%20return%20cls_tokens%2C%20logits%2C%20model_service%2C%20sae_out%0A%0A%0A%40app.cell%0Adef%20_(h_sparse%2C%20img_tensor)%3A%0A%20%20%20%20h_sparse.shape%2C%20img_tensor.max()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(img_tensor%2C%20model_service%2C%20plt)%3A%0A%20%20%20%20logits_svc%20%3D%20model_service.predict(img_tensor)%0A%20%20%20%20plt.imshow(logits_svc%5B0%5D.cpu())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(sae_out)%3A%0A%20%20%20%20sae_out.keys()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(sae_out)%3A%0A%20%20%20%20sae_out%5B'layer_reconstructed'%5D%5B0%5D.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20Opacity%20slider%20control%0A%20%20%20%20opacity_slider%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%20start%3D0.0%2C%0A%20%20%20%20%20%20%20%20stop%3D1.0%2C%0A%20%20%20%20%20%20%20%20step%3D0.05%2C%0A%20%20%20%20%20%20%20%20value%3D0.5%2C%0A%20%20%20%20%20%20%20%20label%3D%22Mask%20Opacity%22%0A%20%20%20%20)%0A%20%20%20%20return%20(opacity_slider%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20%23%20Interactive%20area%20selection%20using%20plotly%0A%20%20%20%20selection_info%20%3D%20mo.ui.text(%0A%20%20%20%20%20%20%20%20value%3D%22No%20selection%20yet%22%2C%0A%20%20%20%20%20%20%20%20label%3D%22Selected%20Area%20Info%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(go%2C%20img_rsz%2C%20logits%2C%20mo%2C%20np%2C%20opacity_slider)%3A%0A%20%20%20%20%23%20Create%20interactive%20plotly%20visualization%20with%20CLICK-based%20selection%0A%0A%20%20%20%20%23%20Prepare%20image%20data%0A%20%20%20%20img_rgb%20%3D%20img_rsz.transpose(%5B1%2C%202%2C%200%5D)%0A%0A%20%20%20%20%23%20Normalize%20to%200-255%20range%20if%20needed%0A%20%20%20%20if%20img_rgb.max()%20%3C%3D%201.0%3A%0A%20%20%20%20%20%20%20%20img_rgb%20%3D%20(img_rgb%20*%20255).astype(np.uint8)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20img_rgb%20%3D%20img_rgb.astype(np.uint8)%0A%0A%20%20%20%20%23%20Create%20mask%20overlay%0A%20%20%20%20mask_data%20%3D%20logits%5B0%5D.argmax(dim%3D0).cpu().numpy()%0A%0A%20%20%20%20%23%20Create%20figure%20-%20using%20scatter%20for%20better%20click%20detection%0A%20%20%20%20fig_interactive%20%3D%20go.Figure()%0A%0A%20%20%20%20%23%20Add%20image%20as%20background%0A%20%20%20%20fig_interactive.add_trace(%0A%20%20%20%20%20%20%20%20go.Image(z%3Dimg_rgb%2C%20name%3D%22image%22)%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20Add%20mask%20as%20heatmap%20overlay%20with%20better%20click%20support%0A%20%20%20%20fig_interactive.add_trace(%0A%20%20%20%20%20%20%20%20go.Heatmap(%0A%20%20%20%20%20%20%20%20%20%20%20%20z%3Dmask_data%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20colorscale%3D'Viridis'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20opacity%3Dopacity_slider.value%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20showscale%3DTrue%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22segmentation%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20hovertemplate%3D'Class%3A%20%25%7Bz%7D%3Cbr%3Ex%3A%20%25%7Bx%7D%3Cbr%3Ey%3A%20%25%7By%7D%3Cextra%3E%3C%2Fextra%3E'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20customdata%3Dmask_data%20%20%23%20Store%20mask%20data%20for%20access%20on%20click%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20)%0A%0A%0A%20%20%20%20%23%20Make%20it%20interactive%20with%20marimo%20-%20captures%20click%20events%0A%20%20%20%20plotly_chart%20%3D%20mo.ui.plotly(fig_interactive)%0A%20%20%20%20return%20img_rgb%2C%20plotly_chart%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20opacity_slider)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%23%20Interactive%20Visualization%20Controls%5Cn%5CnAdjust%20mask%20opacity%3A%20%7Bopacity_slider%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20opacity_slider%2C%20plotly_chart)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%23%20Interactive%20Visualization%0A%0A%20%20%20%20**Instructions%3A**%0A%20%20%20%20-%20Click%20on%20the%20image%20to%20inspect%20specific%20points%0A%20%20%20%20-%20Use%20the%20manual%20selection%20controls%20below%20to%20define%20a%20region%20of%20interest%0A%0A%20%20%20%20%7Bopacity_slider%7D%0A%0A%20%20%20%20%7Bplotly_chart%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(device%2C%20img_size%2C%20plt%2C%20torch)%3A%0A%20%20%20%20x_left%2C%20y_bottom%20%3D%20630%2C%2090%0A%20%20%20%20w_patch%2C%20h_patch%20%3D%20120%2C%2090%0A%20%20%20%20region_mask%20%3D%20torch.zeros(img_size%2C%20img_size).to(device)%0A%20%20%20%20region_mask%5By_bottom%3Ay_bottom%2Bh_patch%2C%20x_left%3Ax_left%2Bw_patch%5D%20%3D%201%0A%20%20%20%20plt.imshow(region_mask.cpu())%0A%20%20%20%20return%20h_patch%2C%20region_mask%2C%20w_patch%2C%20x_left%2C%20y_bottom%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20%23%20Create%20overlay%20visualization%20with%20controlled%20opacity%0A%20%20%20%20%23%20fig%2C%20ax%20%3D%20plt.subplots(figsize%3D(12%2C%208))%0A%0A%20%20%20%20%23%20%23%20Display%20original%20image%0A%20%20%20%20%23%20ax.imshow(img_rsz.transpose(%5B1%2C%202%2C%200%5D))%0A%0A%20%20%20%20%23%20rect%20%3D%20Rectangle((x_left%2C%20y_bottom)%2C%20w_patch%2C%20h_patch%2C%0A%20%20%20%20%23%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20fill%3DFalse%2C%20linewidth%3D2)%20%20%20%0A%20%20%20%20%23%20ax.add_patch(rect)%0A%0A%20%20%20%20%23%20%23%20Create%20mask%20from%20logits%0A%20%20%20%20%23%20labels%20%3D%20logits%5B0%5D.argmax(dim%3D0).cpu().numpy()%0A%0A%20%20%20%20%23%20%23%20Overlay%20mask%20with%20controlled%20opacity%0A%20%20%20%20%23%20ax.imshow(labels%2C%20alpha%3Dopacity_slider.value%2C%20cmap%3D'tab10'%2C%20interpolation%3D'nearest')%0A%0A%20%20%20%20%23%20ax.set_title(f%22Image%20with%20Segmentation%20Overlay%20(opacity%3A%20%7Bopacity_slider.value%3A.2f%7D)%22)%0A%20%20%20%20%23%20ax.axis('off')%0A%20%20%20%20%23%20plt.tight_layout()%0A%0A%20%20%20%20%23%20fig%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h_patch%2C%20img_rgb%2C%20plt%2C%20w_patch%2C%20x_left%2C%20y_bottom)%3A%0A%20%20%20%20img_patch%20%3D%20img_rgb%5By_bottom%3Ay_bottom%2Bh_patch%2C%20x_left%3Ax_left%2Bw_patch%5D%0A%20%20%20%20plt.imshow(img_patch)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Extract%20Features%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(sae_out)%3A%0A%20%20%20%20sae_out.keys()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(sae_out)%3A%0A%20%20%20%20h_sparse%20%3D%20sae_out%5B'h_sparse'%5D%0A%20%20%20%20n_tokens%20%3D%20h_sparse.shape%5B1%5D%0A%20%20%20%20h_token%20%3D%20w_token%20%3D%20int((n_tokens)**0.5)%0A%20%20%20%20sparse_coeffs%20%3D%20h_sparse.reshape(h_sparse.shape%5B0%5D%2C%20h_token%2C%20w_token%2C%20-1)%0A%20%20%20%20h_sparse.shape%2C%20sparse_coeffs.shape%0A%20%20%20%20return%20h_sparse%2C%20h_token%2C%20sparse_coeffs%2C%20w_token%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20h_token%2C%20logits%2C%20w_token)%3A%0A%20%20%20%20mask_patch_shape%20%3D%20F.interpolate(logits.argmax(dim%3D1).unsqueeze(1).float()%2C%20size%20%3D%20(h_token%2C%20w_token)%2C%20mode%20%3D%20'nearest').squeeze(1).cpu()%0A%20%20%20%20mask_patch_shape.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20h_token%2C%20region_mask%2C%20w_token)%3A%0A%20%20%20%20region_mask_patches%20%3D%20F.interpolate(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20region_mask.unsqueeze(0).unsqueeze(0).float()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20size%3D(h_token%2C%20w_token)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mode%3D'nearest'%0A%20%20%20%20%20%20%20%20%20%20%20%20).squeeze().bool()%0A%20%20%20%20region_mask_patches.shape%0A%20%20%20%20return%20(region_mask_patches%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(rf%22%22%22%0A%20%20%20%20You%20get%20the%20number%20of%20patchs%20and%20sparse%20coefficients%20for%20each%20patch%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_mask_patches%2C%20sparse_coeffs%2C%20torch)%3A%0A%20%20%20%20region_features%20%3D%20sparse_coeffs.squeeze()%5Bregion_mask_patches%5D%0A%20%20%20%20mean_features%20%3D%20region_features.mean(dim%3D0)%20%0A%20%20%20%20threshold%20%3D%200.01%20*%20mean_features.max()%0A%20%20%20%20active_mask%20%3D%20mean_features%20%3E%20threshold%0A%20%20%20%20active_indices%20%3D%20active_mask.nonzero().squeeze(-1).tolist()%0A%20%20%20%20region_top_values%2C%20region_top_indices%20%3D%20torch.topk(mean_features%2C%2064)%0A%20%20%20%20region_features.shape%2C%20threshold%2C%20region_top_values%2C%20region_top_indices%0A%20%20%20%20return%20region_features%2C%20region_top_indices%2C%20region_top_values%0A%0A%0A%40app.cell%0Adef%20_(logits)%3A%0A%20%20%20%20logits.argmax(dim%3D1).shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20device%2C%20logits%2C%20sparse_coeffs%2C%20torch)%3A%0A%20%20%20%20%40torch.no_grad()%0A%20%20%20%20def%20compute_class_feature_profiles(%0A%20%20%20%20%20%20%20%20sparse_coeffs%2C%0A%20%20%20%20%20%20%20%20seg_labels%2C%0A%20%20%20%20%20%20%20%20num_classes%3A%20int%20%3D%209%0A%20%20%20%20)%20-%3E%20torch.Tensor%3A%0A%20%20%20%20%20%20%20%20class_profiles%20%3D%20torch.zeros(num_classes%2C%20sparse_coeffs.shape%5B-1%5D%2C%20device%3Ddevice)%0A%20%20%20%20%20%20%20%20class_counts%20%3D%20torch.zeros(num_classes%2C%20device%3Ddevice)%0A%0A%20%20%20%20%20%20%20%20%23%20for%20i%2C%20image%20in%20enumerate(images)%3A%0A%0A%20%20%20%20%20%20%20%20sparse%20%3D%20sparse_coeffs%0A%20%20%20%20%20%20%20%20B%2C%20h%2C%20w%2C%20D%20%3D%20sparse.shape%0A%0A%0A%20%20%20%20%20%20%20%20%23%20Resize%20labels%20to%20feature%20map%20size%0A%20%20%20%20%20%20%20%20if%20seg_labels.shape%5B-2%3A%5D%20!%3D%20(h%2C%20w)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20seg_labels%20%3D%20F.interpolate(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20seg_labels.unsqueeze(1).float()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20size%3D(h%2C%20w)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mode%3D'nearest'%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20).squeeze(1).long()%0A%0A%20%20%20%20%20%20%20%20seg_labels%20%3D%20seg_labels.squeeze(0)%20%20%23%20(h%2C%20w)%0A%20%20%20%20%20%20%20%20sparse_flat%20%3D%20sparse.squeeze(0)%20%20%23%20(h%2C%20w%2C%20d_hidden)%0A%0A%20%20%20%20%20%20%20%20%23%20Accumulate%20for%20each%20class%0A%20%20%20%20%20%20%20%20for%20cls_idx%20in%20range(num_classes)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20mask%20%3D%20(seg_labels%20%3D%3D%20cls_idx)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20mask.any()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20cls_features%20%3D%20sparse_flat%5Bmask%5D%20%20%23%20(N_cls%2C%20d_hidden)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(sparse_flat.shape%2C%20cls_features.shape)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20class_profiles%5Bcls_idx%5D%20%2B%3D%20cls_features.sum(dim%3D0)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20class_counts%5Bcls_idx%5D%20%2B%3D%20mask.sum()%0A%0A%20%20%20%20%20%20%20%20%23%20Average%0A%20%20%20%20%20%20%20%20class_profiles%20%3D%20class_profiles%20%2F%20(class_counts.unsqueeze(1)%20%2B%201e-8)%0A%20%20%20%20%20%20%20%20print(class_counts)%0A%20%20%20%20%20%20%20%20return%20class_profiles%0A%0A%20%20%20%20class_profiles%20%3D%20compute_class_feature_profiles(sparse_coeffs%2C%20logits.argmax(dim%3D1))%0A%20%20%20%20dominant_class%20%3D%20class_profiles.argmax(dim%3D0)%0A%20%20%20%20class_profiles.shape%2C%20dominant_class.shape%0A%20%20%20%20return%20(class_profiles%2C)%0A%0A%0A%40app.cell%0Adef%20_(class_profiles)%3A%0A%20%20%20%20class_profiles.max()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Sparse%20distribution%0A%20%20%20%20The%20scatter%20plot%20shows%20that%20different%20classes%20are%20getting%20activated%20through%20different%20features%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles)%3A%0A%20%20%20%20class_profiles%5B%3A%2C0%3A10000%5D.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles%2C%20px)%3A%0A%20%20%20%20fig_cp%20%3D%20px.imshow(class_profiles%5B%3A%2C0%3A60000%5D.cpu().numpy()%2C%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color_continuous_scale%3D'hot'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20aspect%3D'auto'%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20labels%3D%7B'x'%3A%20'Feature'%2C%20'y'%3A%20'Row'%2C%20'color'%3A%20'Value'%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20zmin%20%3D%200%2C%20zmax%20%3D%203%20)%0A%20%20%20%20fig_cp.update_layout(width%3D1200%2C%20height%3D400)%0A%20%20%20%20fig_cp.write_image(%22tmp%2Ficml%2Fsparse_heatmap.png%22)%0A%20%20%20%20fig_cp.show()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_top_indices%2C%20region_top_values)%3A%0A%20%20%20%20region_top_indices%2C%20region_top_values%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Extract%20class%20region%20features%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles%2C%20region_mask_patches%2C%20sparse_coeffs%2C%20torch)%3A%0A%20%20%20%20def%20get_region_class_profile_match(class_profile%2C%20sparse_coeffs%2C%20region_mask%2C%20top_k%20%3D%2032)%3A%0A%20%20%20%20%20%20%20%20region_features%20%3D%20sparse_coeffs.squeeze()%5Bregion_mask%5D%0A%20%20%20%20%20%20%20%20mean_features%20%3D%20region_features.mean(dim%3D0)%20%0A%20%20%20%20%20%20%20%20threshold%20%3D%200.1%20*%20mean_features.max()%0A%20%20%20%20%20%20%20%20active_mask%20%3D%20mean_features%20%3E%20threshold%0A%20%20%20%20%20%20%20%20region_top_indices%20%3D%20active_mask.nonzero().squeeze(-1)%0A%20%20%20%20%20%20%20%20if%20isinstance(region_top_indices%2C%20int)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20region_top_indices%20%3D%20%5Bregion_top_indices%5D%0A%20%20%20%20%20%20%20%20region_top_values%20%3D%20mean_features%5Bactive_mask%5D%0A%20%20%20%20%20%20%20%20%23%20print(%22Masked%20indices%22)%0A%20%20%20%20%20%20%20%20%23%20print(region_top_indices%2C%20region_top_values)%0A%20%20%20%20%20%20%20%20%23%20region_top_values%2C%20region_top_indices%20%3D%20torch.topk(mean_features%2C%20top_k)%0A%0A%20%20%20%20%20%20%20%20region_class_match_index%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20region_class_match_value%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20idx%20in%20range(len(class_profile))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20class_top_val%2C%20class_top_index%20%3D%20torch.topk(class_profile%5Bidx%5D%2C%20top_k)%0A%20%20%20%20%20%20%20%20%20%20%20%20region_top_class_mask%20%3D%20torch.isin(class_top_index%2C%20region_top_indices)%0A%20%20%20%20%20%20%20%20%20%20%20%20region_class_match_index.append(class_top_index%5Bregion_top_class_mask%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20region_class_match_value.append(class_top_val%5Bregion_top_class_mask%5D)%0A%20%20%20%20%20%20%20%20return%20region_class_match_index%2C%20region_class_match_value%0A%0A%20%20%20%20def%20get_discriminating_features_exclusive(class_profile%2C%20sparse_coeffs%2C%20region_mask%2C%20top_k%3D32)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Find%20class-exclusive%20features%20using%20set%20operations%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20region_features%20%3D%20sparse_coeffs.squeeze()%5Bregion_mask%5D%0A%20%20%20%20%20%20%20%20mean_features%20%3D%20region_features.mean(dim%3D0)%20%0A%20%20%20%20%20%20%20%20region_top_values%2C%20region_top_indices%20%3D%20torch.topk(mean_features%2C%20top_k)%0A%0A%20%20%20%20%20%20%20%20region_class_match_index%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20region_class_match_value%20%3D%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20%23%20Get%20top-k%20for%20all%20classes%0A%20%20%20%20%20%20%20%20all_class_tops%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20idx%20in%20range(len(class_profile))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_%2C%20class_top_index%20%3D%20torch.topk(class_profile%5Bidx%5D%2C%20top_k)%0A%20%20%20%20%20%20%20%20%20%20%20%20all_class_tops.append(set(class_top_index.tolist()))%0A%0A%20%20%20%20%20%20%20%20%23%20For%20each%20class%0A%20%20%20%20%20%20%20%20for%20idx%20in%20range(len(class_profile))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20class_top_val%2C%20class_top_index%20%3D%20torch.topk(class_profile%5Bidx%5D%2C%20top_k)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Region%20features%0A%20%20%20%20%20%20%20%20%20%20%20%20region_set%20%3D%20set(region_top_indices.tolist())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20This%20class%20features%0A%20%20%20%20%20%20%20%20%20%20%20%20class_set%20%3D%20set(class_top_index.tolist())%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20All%20OTHER%20classes%20features%20(union)%0A%20%20%20%20%20%20%20%20%20%20%20%20other_classes_set%20%3D%20set()%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20other_idx%20in%20range(len(class_profile))%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20other_idx%20!%3D%20idx%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20other_classes_set.update(all_class_tops%5Bother_idx%5D)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Exclusive%20features%20%3D%20(region%20%E2%88%A9%20this_class)%20-%20other_classes%0A%20%20%20%20%20%20%20%20%20%20%20%20exclusive_set%20%3D%20(region_set%20%26%20class_set)%20-%20other_classes_set%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Convert%20back%20to%20tensor%20and%20get%20values%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20len(exclusive_set)%20%3E%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_list%20%3D%20list(exclusive_set)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_indices%20%3D%20torch.tensor(exclusive_list%2C%20device%3Dregion_top_indices.device)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Get%20values%20for%20these%20indices%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_values%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20for%20feat_idx%20in%20exclusive_indices%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20Find%20where%20this%20feature%20appears%20in%20region_top_indices%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mask%20%3D%20region_top_indices%20%3D%3D%20feat_idx%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20mask.any()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_values.append(region_top_values%5Bmask%5D%5B0%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_values.append(mean_features%5Bfeat_idx%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_values%20%3D%20torch.tensor(exclusive_values%2C%20device%3Dregion_top_indices.device)%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_indices%20%3D%20torch.tensor(%5B%5D%2C%20device%3Dregion_top_indices.device%2C%20dtype%3Dtorch.long)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20exclusive_values%20%3D%20torch.tensor(%5B%5D%2C%20device%3Dregion_top_indices.device)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20region_class_match_index.append(exclusive_indices)%0A%20%20%20%20%20%20%20%20%20%20%20%20region_class_match_value.append(exclusive_values)%0A%0A%20%20%20%20%20%20%20%20return%20region_class_match_index%2C%20region_class_match_value%0A%0A%20%20%20%20get_region_class_profile_match(class_profiles%2C%20sparse_coeffs%2C%20region_mask_patches)%2C%20get_discriminating_features_exclusive(class_profiles%2C%20sparse_coeffs%2C%20region_mask_patches%2C%2032)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(px%2C%20region_mask_patches)%3A%0A%20%20%20%20px.imshow(region_mask_patches.cpu())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h_sparse)%3A%0A%20%20%20%20(h_sparse%5B0%5D%3E1e-6).shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h_sparse%2C%20torch)%3A%0A%20%20%20%20torch.topk((h_sparse%5B0%5D%3E1e-6).sum(dim%20%3D%200)%2C%2020)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles%2C%20mo%2C%20px%2C%20sparse_coeffs%2C%20torch)%3A%0A%20%20%20%20print(torch.topk(class_profiles%5B4%5D%2C%2015))%0A%20%20%20%20sparse_coeff_plot%20%3D%20mo.ui.plotly(px.imshow(sparse_coeffs%5B0%5D%5B%3A%2C%3A%2C617%5D.cpu()))%0A%20%20%20%20return%20(sparse_coeff_plot%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20sparse_coeff_plot)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20In%20this%20example%20plot%2C%20it%20is%20clearly%20seen%20that%20the%20the%20sparse%20coeeficient%20that%20belongs%20to%20a%20given%20class%20is%20activating%20on%20all%20pixels%20covered%20by%20vegetation.%20Vegetation%20is%20class%205%0A%20%20%20%20%7Bsparse_coeff_plot%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20np%2C%20plt%2C%20sparse_coeffs)%3A%0A%20%20%20%20def%20overlay_image_sparse_coeff(image_rgb%2C%20sparse_coeff%2C%20sparse_idxs%2C%20colormap%3D'viridis'%2C%20alpha%3D0.5%2C%20title%3D%22Overlay%22%2C%20ax%3DNone%2C)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Overlays%20a%20mask%20on%20an%20RGB%20image%20using%20a%20specific%20Matplotlib%20colormap.%0A%20%20%20%20%20%20%20%20Background%20(0)%20values%20in%20the%20mask%20are%20treated%20as%20transparent.%0A%0A%20%20%20%20%20%20%20%20Parameters%3A%0A%20%20%20%20%20%20%20%20-%20image_rgb%3A%20numpy%20array%20of%20shape%20(H%2C%20W%2C%203).%0A%20%20%20%20%20%20%20%20-%20mask%3A%20numpy%20array%20of%20shape%20(H%2C%20W).%20Can%20be%20binary%2C%20integer%20labels%2C%20or%20float%20heatmap.%0A%20%20%20%20%20%20%20%20-%20colormap%3A%20string%2C%20Matplotlib%20colormap%20name%20(e.g.%2C%20'viridis'%2C%20'turbo'%2C%20'plasma').%0A%20%20%20%20%20%20%20%20-%20alpha%3A%20float%20(0%20to%201)%2C%20global%20transparency%20of%20the%20overlay.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20if%20ax%20is%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20fig%2C%20ax%20%3D%20plt.subplots(figsize%3D(8%2C%208))%0A%0A%20%20%20%20%20%20%20%20%23%201.%20Show%20the%20Base%20Image%0A%20%20%20%20%20%20%20%20ax.imshow(image_rgb)%0A%20%20%20%20%20%20%20%20sparse_image%20%3D%20sparse_coeffs%5B%3A%2C%3A%2C%3A%2Csparse_idxs%5B0%5D%5D%0A%20%20%20%20%20%20%20%20for%20sparse_idx%20in%20sparse_idxs%5B1%3A%5D%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20sparse_image%20%3D%20sparse_image%2Bsparse_coeffs%5B%3A%2C%3A%2C%3A%2Csparse_idxs%5B1%5D%5D%0A%0A%20%20%20%20%20%20%20%20mask%20%3D%20F.interpolate(sparse_image.unsqueeze(0)%2Fsparse_image.max()%2C%20size%3D(image_rgb.shape%5B1%5D%2C%20image_rgb.shape%5B1%5D)%2Cmode%3D'nearest').cpu()%5B0%5D%5B0%5D%0A%20%20%20%20%20%20%20%20%23%202.%20Mask%20the%20background%20(where%20mask%20is%200)%0A%20%20%20%20%20%20%20%20%23%20This%20creates%20a%20%22masked%20array%22%20where%200s%20are%20invalid%20and%20won't%20be%20drawn%0A%20%20%20%20%20%20%20%20masked_data%20%3D%20np.ma.masked_where(mask%20%3D%3D%200%2C%20mask)%0A%0A%20%20%20%20%20%20%20%20%23%203.%20Overlay%20the%20Mask%20using%20the%20chosen%20colormap%0A%20%20%20%20%20%20%20%20%23%20'alpha'%20here%20controls%20the%20transparency%20of%20the%20colored%20parts%20only%0A%20%20%20%20%20%20%20%20im%20%3D%20ax.imshow(masked_data%2C%20cmap%3Dcolormap%2C%20alpha%3Dalpha%2C%20vmin%3D0%2C%20vmax%3D1.0)%0A%0A%0A%0A%20%20%20%20%20%20%20%20cbar%20%3D%20plt.colorbar(im%2C%20ax%3Dax%2C%20fraction%3D0.046%2C%20pad%3D0.04)%0A%20%20%20%20%20%20%20%20cbar.set_label('Mask%20Intensity'%2C%20rotation%3D270%2C%20labelpad%3D15)%0A%0A%0A%0A%20%20%20%20%20%20%20%20ax.set_title(title)%0A%20%20%20%20%20%20%20%20ax.axis('off')%0A%20%20%20%20%20%20%20%20return%20im%0A%20%20%20%20return%20(overlay_image_sparse_coeff%2C)%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20img_rgb%2C%20sparse_coeffs)%3A%0A%20%20%20%20F.interpolate(sparse_coeffs%5B%3A%2C%3A%2C%3A%2C617%5D.unsqueeze(0)%2C%20size%3D(1024%2C%201024)%2Cmode%3D'nearest').cpu()%5B0%5D%5B0%5D.shape%2C%20img_rgb.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(img_rgb%2C%20overlay_image_sparse_coeff%2C%20sparse_coeffs)%3A%0A%20%20%20%20overlay_image_sparse_coeff(img_rgb%2C%20sparse_coeffs%2C%20%20%5B59138%5D%2C%20colormap%3D'viridis')%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(img_rgb%2C%20overlay_image_sparse_coeff%2C%20plt%2C%20sparse_coeffs)%3A%0A%20%20%20%20fig%2C%20axes%20%3D%20plt.subplots(1%2C%204%2C%20figsize%3D(12%2C%205))%0A%20%20%20%20axes%5B0%5D.imshow(img_rgb)%0A%20%20%20%20axes%5B0%5D.set_title(%22Image%22)%0A%20%20%20%20axes%5B0%5D.axis('off')%0A%20%20%20%20overlay_image_sparse_coeff(img_rgb%2C%20sparse_coeffs%2C%20%20%5B13128%2C%2014668%5D%2C%20colormap%3D'viridis'%2C%20ax%20%3D%20axes%5B1%5D%2C%20title%20%3D%20%22Tree%20Activations%22)%0A%20%20%20%20overlay_image_sparse_coeff(img_rgb%2C%20sparse_coeffs%2C%20%20%5B38331%2C%2023403%5D%2C%20colormap%3D'viridis'%2C%20ax%20%3D%20axes%5B2%5D%2C%20title%3D%22Road%20Activations%22)%0A%20%20%20%20overlay_image_sparse_coeff(img_rgb%2C%20sparse_coeffs%2C%20%20%5B5368%2C7230%5D%2C%20colormap%3D'viridis'%2C%20ax%20%3D%20axes%5B3%5D%2C%20title%3D%22Building%20activations%22)%0A%20%20%20%20plt.savefig(%22tmp%2Ficml%2Factivation_map.png%22%2C%20dpi%3D300%2C%20bbox_inches%3D'tight')%0A%20%20%20%20return%0A%0A%0Aapp._unparsable_cell(%0A%20%20%20%20r%22%22%22%0A%20%20%20%20%20asprint(torch.topk(class_profiles%5B8%5D%2C%2015))%0A%20%20%20%20%22%22%22%2C%0A%20%20%20%20name%3D%22_%22%0A)%0A%0A%0A%40app.cell%0Adef%20_(img_rgb%2C%20overlay_image_sparse_coeff%2C%20sparse_coeffs)%3A%0A%20%20%20%20overlay_image_sparse_coeff(img_rgb%2C%20sparse_coeffs%2C%20%20%5B5368%2C7230%5D%2C%20colormap%3D'viridis')%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Steer%20Features%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(cls_tokens%2C%20h_sparse%2C%20model_service)%3A%0A%20%20%20%20model_service.get_output_from_sparse_ceoff(h_sparse%2C%20cls_tokens).shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_mask_patches%2C%20sparse_coeffs)%3A%0A%20%20%20%20sparse_coeffs.shape%2C%20region_mask_patches.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles%2C%20torch)%3A%0A%20%20%20%20class_profiles%5B4%5D%0A%20%20%20%20torch.topk(class_profiles%5B4%5D%2C%205)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%22tensor(%5B%209441%2C%20%202096%2C%2010935%2C%2014174%2C%20%206521%2C%20%202693%2C%20%201138%2C%20%207021%2C%20%20%20360%2C%20%208829%2C%5Cn%20%20%20%20%20%20%20%20%202810%2C%20%203007%2C%2011139%2C%20%205623%2C%2012214%2C%20%209582%2C%20%202908%2C%20%209166%2C%20%208082%2C%20%206490%2C%5Cn%20%20%20%20%20%20%20%2013970%2C%20%206829%2C%20%204576%5D%2C%20device%3D'cuda%3A1')%22%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_mask_patches%2C%20sparse_coeffs%2C%20torch)%3A%0A%20%20%20%20sparse_coeffs_cp%20%3D%20sparse_coeffs.clone()%0A%20%20%20%20sparse_coeffs_flat%20%3D%20sparse_coeffs_cp.view(1%2C%20-1%2C%2016384)%0A%20%20%20%20%23%20feature_indices_target%20%3D%20%5B12391%2C%20908%5D%0A%20%20%20%20feature_indices_target%20%3D%20%5B8082%2C%20452%5D%0A%20%20%20%20%23%20feature_indices_supress%20%3D%20%5B8675%2C14634%2C%2010347%2C%20%202768%2C%2011253%2C%2011213%2C%2012231%2C%20%201597%2C%2011771%2C%2012512%5D%0A%20%20%20%20feature_indices_supress%20%3D%20%5B%202693%5D%0A%20%20%20%20for%20feat_idx%20in%20feature_indices_target%3A%0A%20%20%20%20%20%20%20%20sparse_coeffs_flat%5B0%2C%20region_mask_patches.view(-1)%2C%20feat_idx%5D%20%2B%3D%20torch.tensor(5000.0)%0A%20%20%20%20%23%20sparse_coeffs_cp%5B0%2C%20region_mask_patches%5B%3A%2C%20None%5D%2C%20feature_indices_target%5D%20%3D%204.0%0A%20%20%20%20%23%20sparse_coeffs_cp%5B0%2C%20region_mask_patches%2C%20%3A%5D%5B%3A%2C%20feature_indices_supress%5D%20%3D%200.2%0A%20%20%20%20for%20feat_idx%20in%20feature_indices_supress%3A%0A%20%20%20%20%20%20%20%20sparse_coeffs_flat%5B0%2C%20region_mask_patches.view(-1)%2C%20feat_idx%5D%20%3D%200.00%0A%20%20%20%20%23%20sparse_coeffs_flat%20%3D%20%20torch.zeros(1%2C4096%2C%2016384%2C%20device%3D%20device)%0A%20%20%20%20return%20(sparse_coeffs_flat%2C)%0A%0A%0A%40app.cell%0Adef%20_(px%2C%20sparse_coeffs_flat)%3A%0A%20%20%20%20px.imshow(sparse_coeffs_flat%5B0%5D%5B%3A%2C8082%5D.view(64%2C64).cpu())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_mask_patches%2C%20sparse_coeffs_flat)%3A%0A%20%20%20%20sparse_coeffs_flat%5B0%2C%20region_mask_patches.view(-1)%2C%2014634%5D%2C%20sparse_coeffs_flat.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(cls_tokens%2C%20model_service%2C%20px%2C%20sparse_coeffs_flat)%3A%0A%20%20%20%20px.imshow(model_service.get_output_from_sparse_ceoff(sparse_coeffs_flat%2C%20cls_tokens).cpu()%5B0%5D)%20%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Gap%20Filling%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20F%2C%0A%20%20%20%20class_profiles%2C%0A%20%20%20%20cls_tokens%2C%0A%20%20%20%20model_service%2C%0A%20%20%20%20px%2C%0A%20%20%20%20region_features%2C%0A%20%20%20%20region_mask_patches%2C%0A%20%20%20%20sparse_coeffs%2C%0A)%3A%0A%20%20%20%20target_class_idx%20%3D%205%0A%20%20%20%20target_profile%20%3D%20class_profiles%5Btarget_class_idx%5D%0A%20%20%20%20region_mean%20%3D%20region_features.mean(dim%20%3D%200)%0A%20%20%20%20gap%20%3D%20F.relu(target_profile%20-%20region_mean)%0A%20%20%20%20sparse_coeffs_gap%20%3D%20sparse_coeffs.clone().view(1%2C%20-1%2C%20sparse_coeffs.shape%5B-1%5D)%0A%20%20%20%20sparse_coeffs_gap%5B0%2C%20region_mask_patches.view(-1)%5D%20%2B%3D%2010*gap%0A%20%20%20%20px.imshow(model_service.get_output_from_sparse_ceoff(sparse_coeffs_gap%2C%20cls_tokens).cpu()%5B0%5D)%20%0A%20%20%20%20return%20region_mean%2C%20sparse_coeffs_gap%2C%20target_class_idx%2C%20target_profile%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20cls_tokens%2C%0A%20%20%20%20go%2C%0A%20%20%20%20h_patch%2C%0A%20%20%20%20img_rgb%2C%0A%20%20%20%20img_tensor%2C%0A%20%20%20%20make_subplots%2C%0A%20%20%20%20model_service%2C%0A%20%20%20%20sparse_coeffs_gap%2C%0A%20%20%20%20w_patch%2C%0A%20%20%20%20x_left%2C%0A%20%20%20%20y_bottom%2C%0A)%3A%0A%20%20%20%20fig_gap_filling%20%3D%20make_subplots(%0A%20%20%20%20%20%20%20%20rows%3D1%2C%20cols%3D3%2C%0A%20%20%20%20%20%20%20%20subplot_titles%3D(%22Base%20Image%22%2C%20%22Model%20prediction%22%2C%20%22Steered%20Output%22)%2C%0A%20%20%20%20%20%20%20%20%23%20horizontal_spacing%3D0.1%0A%20%20%20%20)%0A%0A%20%20%20%20%23%20---%203.%20Add%20Traces%20---%0A%0A%20%20%20%20%23%20Left%20Plot%20(Row%201%2C%20Col%201)%3A%20Just%20the%20Base%20Image%0A%20%20%20%20fig_gap_filling.add_trace(%0A%20%20%20%20%20%20%20%20go.Image(z%3Dimg_rgb)%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%20col%3D1%0A%20%20%20%20)%0A%20%20%20%20fig_gap_filling.add_trace(%0A%20%20%20%20%20%20%20%20go.Heatmap(z%3Dmodel_service.predict(img_tensor).cpu()%5B0%5D.numpy()%2C%20colorscale%3D'viridis')%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%20col%3D2%0A%20%20%20%20)%0A%20%20%20%20fig_gap_filling.update_yaxes(autorange%3D%22reversed%22%2C%20row%3D1%2C%20col%3D2)%0A%0A%0A%20%20%20%20%23%20Right%20Plot%20(Row%201%2C%20Col%202)%3A%20Base%20Image...%0A%20%20%20%20fig_gap_filling.add_trace(%0A%20%20%20%20%20%20%20%20go.Heatmap(z%3Dmodel_service.get_output_from_sparse_ceoff(sparse_coeffs_gap%2C%20cls_tokens).cpu()%5B0%5D.numpy()%2C%20colorscale%3D'viridis')%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%20col%3D3%0A%20%20%20%20)%0A%20%20%20%20fig_gap_filling.update_yaxes(autorange%3D%22reversed%22%2C%20row%3D1%2C%20col%3D3)%0A%0A%20%20%20%20fig_gap_filling.add_shape(%0A%20%20%20%20%20%20%20%20type%3D%22rect%22%2C%0A%20%20%20%20%20%20%20%20x0%3Dx_left%2C%20y0%3Dy_bottom%2Bh_patch%2C%20%20%23%20Bottom-left%20corner%20(or%20Top-left%20for%20images)%0A%20%20%20%20%20%20%20%20x1%3Dx_left%2Bw_patch%2C%20y1%3Dy_bottom%2C%20%20%23%20Top-right%20corner%20(or%20Bottom-right%20for%20images)%0A%20%20%20%20%20%20%20%20line%3Ddict(color%3D%22Red%22%2C%20width%3D3)%2C%0A%20%20%20%20%20%20%20%20fillcolor%3D%22rgba(255%2C%200%2C%200%2C%200.0)%22%2C%20%23%20Optional%20transparent%20fill%2C%0A%20%20%20%20%20%20%20%20row%3D1%2C%20col%3D1%0A%20%20%20%20)%0A%20%20%20%20fig_gap_filling.update_layout(%0A%20%20%20%20%20%20%20%20width%3D1200%2C%20%20%20%23%20Total%20width%20in%20pixels%0A%20%20%20%20%20%20%20%20%23%20height%3D500%2C%20%20%20%23%20Total%20height%20in%20pixels%0A%20%20%20%20%20%20%20%20autosize%3DTrue%2C%20%23%20Prevents%20the%20browser%20from%20overriding%20your%20size%0A%20%20%20%20%20%20%20%20title%3D%22Steering%20the%20model%20output%20towards%20desired%20class%20(from%20bare%20to%20tree%20class)%22%0A%20%20%20%20)%0A%20%20%20%20fig_gap_filling.write_image(%22tmp%2Ficml%2Fpatch_steering.pdf%22)%0A%20%20%20%20fig_gap_filling%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20F%2C%0A%20%20%20%20cls_tokens%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20model_service%2C%0A%20%20%20%20px%2C%0A%20%20%20%20region_mask_patches%2C%0A%20%20%20%20sparse_coeffs%2C%0A%20%20%20%20target_profile%2C%0A%20%20%20%20torch%2C%0A)%3A%0A%20%20%20%20%40torch.no_grad()%0A%20%20%20%20def%20propagate_steering_to_similar_regions(%0A%20%20%20%20%20%20%20%20sparse_coeffs%3A%20torch.Tensor%2C%20%20%20%20%20%20%23%20(1%2C%20h%2C%20w%2C%20d_hidden)%0A%20%20%20%20%20%20%20%20region_mask%3A%20torch.Tensor%2C%20%20%20%20%20%20%20%20%20%23%20(h%2C%20w)%20bool%20-%20your%20selected%20region%0A%20%20%20%20%20%20%20%20target_class_profile%3A%20torch.Tensor%2C%20%23%20(d_hidden%2C)%20-%20target%20class%20features%0A%20%20%20%20%20%20%20%20similarity_threshold%3A%20float%20%3D%200.5%2C%20%20%23%20cosine%20similarity%20threshold%0A%20%20%20%20%20%20%20%20top_k_features%3A%20int%20%3D%2064%2C%20%20%20%20%20%20%20%20%20%20%23%20number%20of%20features%20to%20consider%0A%20%20%20%20)%3A%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20Find%20all%20pixels%20similar%20to%20selected%20region%20and%20apply%20gap-filling.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20h%2C%20w%2C%20d_hidden%20%3D%20sparse_coeffs.shape%5B1%5D%2C%20sparse_coeffs.shape%5B2%5D%2C%20sparse_coeffs.shape%5B3%5D%0A%20%20%20%20%20%20%20%20sparse_flat%20%3D%20sparse_coeffs.squeeze(0).reshape(-1%2C%20d_hidden)%20%20%23%20(h*w%2C%20d_hidden)%0A%0A%20%20%20%20%20%20%20%20%23%20Step%201%3A%20Get%20feature%20signature%20of%20selected%20region%0A%20%20%20%20%20%20%20%20region_mask_flat%20%3D%20region_mask.view(-1)%20%20%23%20(h*w%2C)%0A%20%20%20%20%20%20%20%20region_features%20%3D%20sparse_flat%5Bregion_mask_flat%5D%20%20%23%20(N_region%2C%20d_hidden)%0A%20%20%20%20%20%20%20%20region_signature%20%3D%20region_features.mean(dim%3D0)%20%20%20%23%20(d_hidden%2C)%20-%20prototype%0A%0A%20%20%20%20%20%20%20%20%23%20Step%202%3A%20Find%20which%20features%20define%20this%20region%20(top-k%20active)%0A%20%20%20%20%20%20%20%20top_vals%2C%20top_indices%20%3D%20torch.topk(region_signature%2C%20top_k_features)%0A%0A%20%20%20%20%20%20%20%20%23%20Step%203%3A%20Compute%20similarity%20of%20ALL%20pixels%20to%20region%20signature%0A%20%20%20%20%20%20%20%20%23%20Only%20compare%20on%20the%20active%20features%20(avoids%20zero-dilution)%0A%20%20%20%20%20%20%20%20region_active%20%3D%20region_signature%5Btop_indices%5D%20%20%23%20(top_k%2C)%0A%20%20%20%20%20%20%20%20all_pixels_active%20%3D%20sparse_flat%5B%3A%2C%20top_indices%5D%20%20%23%20(h*w%2C%20top_k)%0A%0A%20%20%20%20%20%20%20%20%23%20Cosine%20similarity%20between%20each%20pixel%20and%20region%20prototype%0A%20%20%20%20%20%20%20%20similarity%20%3D%20F.cosine_similarity(%0A%20%20%20%20%20%20%20%20%20%20%20%20all_pixels_active%2C%20%0A%20%20%20%20%20%20%20%20%20%20%20%20region_active.unsqueeze(0)%2C%20%20%23%20(1%2C%20top_k)%0A%20%20%20%20%20%20%20%20%20%20%20%20dim%3D-1%0A%20%20%20%20%20%20%20%20)%20%20%23%20(h*w%2C)%0A%0A%20%20%20%20%20%20%20%20%23%20Step%204%3A%20Create%20mask%20of%20similar%20pixels%0A%20%20%20%20%20%20%20%20similar_mask%20%3D%20similarity%20%3E%20similarity_threshold%20%20%23%20(h*w%2C)%0A%0A%20%20%20%20%20%20%20%20%23%20Step%205%3A%20Apply%20gap-filling%20ONLY%20to%20similar%20pixels%0A%20%20%20%20%20%20%20%20gap%20%3D%20F.relu(target_class_profile%20-%20region_signature)%20%20%23%20(d_hidden%2C)%0A%0A%20%20%20%20%20%20%20%20%23%20Clone%20and%20modify%0A%20%20%20%20%20%20%20%20sparse_modified%20%3D%20sparse_flat.clone()%0A%20%20%20%20%20%20%20%20sparse_modified%5Bsimilar_mask%5D%20%3D%20sparse_modified%5Bsimilar_mask%5D%20%2B%205*gap%0A%0A%20%20%20%20%20%20%20%20%23%20Reshape%20back%0A%20%20%20%20%20%20%20%20%23%20sparse_modified%20%3D%20sparse_modified.reshape(1%2C%20h%2C%20w%2C%20d_hidden)%0A%0A%20%20%20%20%20%20%20%20return%20sparse_modified.unsqueeze(0)%2C%20similar_mask.reshape(h%2C%20w)%0A%0A%20%20%20%20sparse_coeffs_all_region%2C%20similar_mask%20%3D%20propagate_steering_to_similar_regions(sparse_coeffs%2C%20region_mask_patches%2C%20target_profile)%0A%0A%20%20%20%20steered_output%20%3D%20model_service.get_output_from_sparse_ceoff(sparse_coeffs_all_region%2C%20cls_tokens).cpu()%5B0%5D%0A%0A%20%20%20%20global_steering%20%3D%20px.imshow(steered_output%2C%20color_continuous_scale%3D'viridis')%0A%20%20%20%20%23%20global_steering.write_image(%22tmp%2Ficml%2Fglobal_steering.pdf%22)%0A%20%20%20%20gap_steering_fig%20%3D%20mo.ui.plotly(global_steering)%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20gap_steering_fig%2C%0A%20%20%20%20%20%20%20%20similar_mask%2C%0A%20%20%20%20%20%20%20%20sparse_coeffs_all_region%2C%0A%20%20%20%20%20%20%20%20steered_output%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(gap_steering_fig%2C%20mo)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20The%20area%20belonging%20to%20bareland%20was%20successfully%20converted%20into%20forest%20area%0A%20%20%20%20%7Bgap_steering_fig%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20img_rgb%2C%20np%2C%20plt%2C%20similar_mask%2C%20steered_output)%3A%0A%20%20%20%20def%20plot_global_steering(image%2C%20similar_mask%2C%20steered_output%2C%20alpha%20%3D%200.6)%3A%0A%20%20%20%20%20%20%20%20fig%2C%20axs%20%3D%20plt.subplots(1%2C2%2C%20figsize%3D(10%2C%205))%0A%20%20%20%20%20%20%20%20axs%5B0%5D.imshow(image)%0A%0A%20%20%20%20%20%20%20%20mask%20%3D%20F.interpolate(similar_mask.float().unsqueeze(0).unsqueeze(0)%2C%20size%3D(1024%2C1024)%2Cmode%3D'nearest').squeeze().cpu()%0A%20%20%20%20%20%20%20%20%23%202.%20Mask%20the%20background%20(where%20mask%20is%200)%0A%20%20%20%20%20%20%20%20%23%20This%20creates%20a%20%22masked%20array%22%20where%200s%20are%20invalid%20and%20won't%20be%20drawn%0A%20%20%20%20%20%20%20%20masked_data%20%3D%20np.ma.masked_where(mask%20%3D%3D%200%2C%20mask)%0A%20%20%20%20%20%20%20%20print(image.shape%2C%20masked_data.shape)%0A%20%20%20%20%20%20%20%20%23%203.%20Overlay%20the%20Mask%20using%20the%20chosen%20colormap%0A%20%20%20%20%20%20%20%20%23%20'alpha'%20here%20controls%20the%20transparency%20of%20the%20colored%20parts%20only%0A%20%20%20%20%20%20%20%20axs%5B0%5D.imshow(masked_data%2C%20cmap%20%3D%20'viridis'%2C%20alpha%3Dalpha%2C%20vmin%3D0.9%2C%20vmax%3D2.0)%0A%20%20%20%20%20%20%20%20axs%5B0%5D.set_title(%22Image%20patches%20similar%20to%20selected%20region%22)%0A%20%20%20%20%20%20%20%20%23%20axs%5B0%5D.axis('off')%0A%20%20%20%20%20%20%20%20axs%5B1%5D.imshow(steered_output)%0A%20%20%20%20%20%20%20%20axs%5B1%5D.set_title(%22Steered%20Output%22)%0A%20%20%20%20%20%20%20%20return%20fig%0A%0A%20%20%20%20steered_figure%20%3D%20plot_global_steering(img_rgb%2C%20similar_mask%2C%20steered_output)%0A%20%20%20%20steered_figure.savefig(%22tmp%2Ficml%2Fsteered_output.png%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(sparse_coeffs_all_region)%3A%0A%20%20%20%20sparse_coeffs_all_region.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(img_tensor%2C%20model_service%2C%20plt)%3A%0A%20%20%20%20plt.imshow(model_service.predict(img_tensor).cpu()%5B0%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%23%20Steer%20Decoder%20Features%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles%2C%20torch)%3A%0A%20%20%20%20torch.topk(class_profiles%5B5%5D%2C%205)%2C%20torch.topk(class_profiles%5B2%5D%2C%205)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%22tensor(%5B%206260%2C%20%202429%2C%20%208282%2C%20%203115%2C%2013877%2C%2010347%2C%20%20%20908%2C%20%208990%2C%2015511%2C%20%20%20893%2C10625%2C%20%202768%2C%20%204964%2C%2014892%2C%2012794%2C%20%205288%2C%20%203871%2C%20%208304%2C%20%208595%2C%20%20%20811%2C%2015484%2C%20%20%20%2055%2C%20%203307%2C%20%205723%2C%2015591%2C%20%204674%2C%20%205899%2C%2014024%2C%20%203856%2C%2012855%2C%5Cn%20%20%20%20%20%20%20%20%202701%2C%20%201824%5D%2C%20device%3D'cuda%3A1')%22%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_features%2C%20torch)%3A%0A%20%20%20%20torch.topk((region_features%20%3E%201e-6).sum(dim%3D0)%2C%2032)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles)%3A%0A%20%20%20%20class_profiles%5B%5B0%2C1%2C2%2C3%2C4%2C6%2C7%2C8%5D%5D.sum().shape%2C%20class_profiles.shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(class_profiles%2C%20region_features%2C%20target_class_idx%2C%20torch)%3A%0A%0A%20%20%20%20target_top_val%2C%20target_top_idx%20%3D%20torch.topk(class_profiles%5B5%5D%2C32)%0A%20%20%20%20other_top_val%2C%20other_top_idx%20%3D%20%5B%5D%2C%20%5B%5D%0A%20%20%20%20for%20idx%20in%20range(len(class_profiles))%3A%0A%20%20%20%20%20%20%20%20if%20idx%20%3D%3D%20target_class_idx%3A%20%20%23%20%E2%86%90%20Skip%20the%20target%20class!%0A%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%20%20%20%20%20%20%20%20class_top_val%2C%20class_top_idx%20%3D%20torch.topk(class_profiles%5Bidx%5D%2C32)%0A%20%20%20%20%20%20%20%20other_top_val.append(class_top_val)%0A%20%20%20%20%20%20%20%20other_top_idx.append(class_top_idx)%0A%0A%20%20%20%20other_top_val%20%3D%20torch.cat(other_top_val)%0A%20%20%20%20other_top_idx%20%3D%20torch.cat(other_top_idx)%0A%20%20%20%20region_top_val%2C%20region_top_idx%20%3D%20torch.topk(region_features.max(dim%20%3D%200)%5B0%5D%2C%2032)%0A%0A%0A%20%20%20%20target_exclusive_set%20%3D%20((set(target_top_idx.cpu().tolist())%20-%20set(other_top_idx.cpu().tolist())))%0A%20%20%20%20target_exclusive%20%3D%20%5Bi%20for%20i%20in%20target_top_idx.cpu().tolist()%20if%20i%20in%20target_exclusive_set%5D%0A%20%20%20%20return%20target_exclusive%2C%20target_top_idx%0A%0A%0A%40app.cell%0Adef%20_(target_exclusive)%3A%0A%20%20%20%20target_exclusive%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(target_top_idx)%3A%0A%20%20%20%20(target_top_idx)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(region_mean%2C%20target_profile%2C%20torch)%3A%0A%20%20%20%20set(torch.topk(region_mean%2C%2032)%5B1%5D.cpu().tolist())%20-%20set(torch.topk(target_profile%2C%2032)%5B1%5D.cpu().tolist())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(img_tensor%2C%20model_service%2C%20plt%2C%20target_exclusive)%3A%0A%20%20%20%20w_dec%20%3D%20model_service.W_dec.clone()%0A%20%20%20%20b_dec%20%3D%20model_service.b_dec.clone()%0A%20%20%20%20w_dec.shape%2C%20b_dec.shape%0A%0A%20%20%20%20features_to_ablate%20%3D%20%5B17664%2C%2033665%2C%2059138%2C%2024196%2C%2035847%2C%2051465%2C%2027534%2C%2040347%2C%2055836%2C%2022941%2C%209506%2C%2016295%2C%2036903%2C%2037041%2C%2052916%2C%2047668%2C%207866%2C%2026435%2C%2011214%2C%2049103%2C%2013648%2C%2013153%2C%2056676%2C%2015462%2C%209581%2C%209453%2C%2062705%2C%2054643%2C%207414%2C%2031991%2C%2023803%2C%2061181%5D%0A%20%20%20%20%23%20w_dec%5B%3A%2C6260%5D%20%3D%20w_dec%5B%3A%2C13091%5D%0A%20%20%20%20w_dec%5B%3A%2C%20features_to_ablate%5D%20%3D%206.0*w_dec%5B%3A%2Ctarget_exclusive%5D.mean(dim%3D1).unsqueeze(1)%0A%0A%20%20%20%20%23%20w_dec%5B%3A%2C%20features_to_ablate%5D%20%3D%200%0A%0A%20%20%20%20%23%20b_dec%5B6260%5D%20%3D%20b_dec%5B3091%5D%0A%20%20%20%20(w_dec%5B%3A%2C6260%5D%20-%20w_dec%5B%3A%2C13091%5D).abs().sum()%0A%20%20%20%20plt.imshow(model_service.predict_custom_features(img_tensor%2C%20w_dec%2C%20b_dec).cpu()%5B0%5D)%0A%20%20%20%20plt.axis(%22off%22)%0A%20%20%20%20plt.savefig(%22tmp%2Ficml%2Fglobal_steering.png%22)%0A%20%20%20%20return%20(w_dec%2C)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(target_exclusive%2C%20w_dec)%3A%0A%20%20%20%20w_dec%5B%3A%2Ctarget_exclusive%5D.mean(dim%3D1).shape%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(w_dec)%3A%0A%20%20%20%20w_dec%5B%3A%2C2429%5D%2C%20w_dec%5B%3A%2C13091%5D%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A